#!/usr/bin/env python
# -*- coding: utf-8 -*-
import matplotlib.pyplot as plt
import numpy as np
import statsmodels.api as sm
import pytz
import sys

from datetime import datetime,timedelta
from zipline.algorithm import TradingAlgorithm
from zipline.transforms import batch_transform
from zipline.utils.factory import load_from_yahoo
from zipline.finance.commission import PerTrade, PerShare, PerDollar
from zipline.finance.slippage import FixedSlippage
from zipline.transforms import MovingAverage, MovingStandardDev
from matplotlib.widgets import Cursor

import matplotlib.backends.backend_pdf

ymd=datetime.now().strftime("%Y-%m-%d")


#outputdir='C:\\_src\\twintower\\twintower\\analysis'
outputdir='C:\\Users\\Public\\_tt\\backtest'
seedcap = 2000

class Pairtrade(TradingAlgorithm):
    def initialize(self, stocks, thresholds=(0.4,0.5,0.9,0.8), basic_unit=20, mavg_window_len=22*2): #如果两线经常交替则不能选得太大，如果平稳则可以选大点(x)
        self.spreads = []
        self.invested = 0
        #self.set_commission(PerShare(0.005))
        self.set_commission(PerTrade(1.1))
        #self.set_commission(PerDollar(0.005))

        self.instant_fill=True
        #self.blotter=blt
        self.set_slippage(slippage=FixedSlippage(spread=0))
        self.capital_base = seedcap
        self.zscores = []
        self.stocks = stocks
        self.unit = basic_unit
        self.window_len=mavg_window_len
        self.add_transform(MovingAverage, 'mavg', ['price'], window_length=mavg_window_len)
        self.add_transform(MovingStandardDev, 'mstd', window_length=mavg_window_len)
        self.daycount=0
        self.basic_unit=basic_unit
        """self.liquidated=False
        self.sell=False
        self.buy=False"""
        self.thresholds = thresholds
        self.tradedaysincelastliquidation = 0 # trade times since last liquidation
        self.cost = 0
        self.stoprisk = False
        self.CumulativePNLSinceLastLiquidation = 0

    def handle_data(self, data):
        self.liquidated=False
        self.sell=False
        self.buy=False
        self.daycount += 1
        self.mavg = data['spread'].mavg['price']
        self.mstd = 0
        if data['spread'].mstd:
            self.mstd = data['spread'].mstd
        spread=data['spread']['price']

        #if self.stocks[0]=='ATHM':
        #    self.mavg = 0
        liq_upper=self.mavg+(self.mstd*self.thresholds[0]) # upper
        liq_lower=self.mavg-(self.mstd*self.thresholds[1]) # lower
        sellspread_threshold=self.mavg+(self.mstd*self.thresholds[2])
        buyspread_threshold=self.mavg-(self.mstd*self.thresholds[3])
        #be careful: self._portfolio['pnl'] is cumulative value
        self.CumulativePNLSinceLastLiquidation = (0 if not self._portfolio else self._portfolio['pnl'])
        #print (0 if not self._portfolio else self._portfolio['pnl'])

        #do not trade when training data
        if self.daycount > self.window_len:
            m=self.mavg
            vol=self.mstd
            #freq = 12
            #self.unit = abs(self.capital_base / self.mavg / (freq) if self.mavg>0.5 else self.capital_base/freq);

            #>>> self._portfolio
            #Out[2]: Portfolio({'portfolio_value': 3000.0, 'cash': 3000.0, 'starting_cash': 3000.0, 'returns': 0.0, 'capital_used': 0.0, 'pnl': 0.0, 'positions': {}, 'positions_value': 0.0, 'start_date': Timestamp('2013-12-11 00:00:00+0000', tz='UTC')})
            if self._portfolio['cash'] < 500 or self._portfolio['portfolio_value'] < 500: # reduce drawdown
                self.stoprisk = True
            #m=0
            #vol=1
            if (liq_lower<spread and spread< liq_upper) and self.invested:# and self.CumulativePNLSinceLastLiquidation>0:
                self.liquidate()
                self.invested = False
                self.liquidated=True
                self.tradedaysincelastliquidation = 0
                self.unit=self.basic_unit
                self.cost = 0;
                self.stoprisk = False
                self.CumulativePNLSinceLastLiquidation = 0
            elif spread>sellspread_threshold and not self.stoprisk:# and not self.invested: # short spread
                self.order(self.stocks[0], -self.unit)
                self.order(self.stocks[1], +self.unit)
                self.invested = True
                self.sell=True
                self.tradedaysincelastliquidation +=1
                self.cost -= spread * self.unit
            elif spread<buyspread_threshold and not self.stoprisk:# and not self.invested: #long spread
                self.order(self.stocks[0], +self.unit)
                self.order(self.stocks[1], -self.unit)
                self.invested = True
                self.buy=True
                self.tradedaysincelastliquidation +=1
                self.cost += spread * self.unit

        self.record(mavg=self.mavg,
                 mstd=self.mstd,
                 spread=spread,
                 liq_upper=liq_upper,
                 liq_lower=liq_lower,
                 sellspread_threshold=sellspread_threshold,
                 buyspread_threshold=buyspread_threshold,
                 liquidated=self.liquidated,
                 tradedaysincelastliquidation=self.tradedaysincelastliquidation,
                 cost=self.cost,
                 sell=self.sell,
                 buy=self.buy,
                 stock0=data[self.stocks[0]]['price']
                 )

    def liquidate(self):
        _BITA_amount = self.portfolio.positions[self.stocks[1]].amount
        self.order(self.stocks[1], -1 * _BITA_amount)
        _ATHM_amount = self.portfolio.positions[self.stocks[0]].amount
        self.order(self.stocks[0], -1 * _ATHM_amount)

def pairplot(results, data, _stocks, totalreturn=0,maxdrawdown=0,thresholds=()):
    fig = plt.figure(figsize=(24.0, 20.0))
    ax1 = plt.subplot(221)
    plt.title("total return:{}%, maxdrawdown={}".format(totalreturn*100,maxdrawdown))
    plt.setp(ax1.get_xticklabels(), visible=True)
    data[_stocks].plot(ax=ax1,style='o-')

    #ax1.legend().set_visible(False)
    ax1.plot(results.ix[results.buy].index, results.stock0[results.buy],'^', markersize=5, color='black')
    ax1.plot(results.ix[results.sell].index, results.stock0[results.sell],'v', markersize=5, color='w')
    ax1.plot(results.ix[results.liquidated].index, results.stock0[results.liquidated],'8', markersize=15, color='r')

    cursor1 = Cursor(ax1, useblit=True, color='red', linewidth=1)
    #ax1.grid(True, color='#fdbdbd')#, alpha=1, lw=1)
    #ax1.get_xaxis().set_visible(False) # hide x axis
    #plt.ylabel('Stock Price')

    ax2 = plt.subplot(222, sharex=ax1)
    plt.title("thresholds:{}".format(thresholds))
    results.portfolio_value.plot(ax=ax2, lw=1, style='o-',legend=True) #mark to market
    results.ending_cash.plot(ax=ax2, lw=1, style='--',color='r',legend=True) # cash in hand
    ax2.plot(results.ix[results.liquidated].index, results.portfolio_value[results.liquidated],'8', markersize=10, color='r')
    #ax2.get_xaxis().set_visible(False) # hide x axis
    #ax2.grid(True, color='#fdbdbd')#, alpha=1, lw=1)
    plt.setp(ax2.get_xticklabels(), visible=True)
    cursor2 = Cursor(ax2, useblit=True, color='blue', linewidth=1 )
    #ax2.set_title('Portfolio, $\n', color="w")
    #results.zscores.plot(ax=ax2, color='r')

    #plt.ylabel('Portfolio($)', color="r")
    #plt.gcf().set_size_inches(18, 8)

    ax4 = plt.subplot(223, sharex=ax1)
    #data['spread'].plot(style='o-',marker=r'$\clubsuit$')
    #data['spread'].plot(style='o-')
    plt.setp(ax4.get_xticklabels(), visible=True)
    ax=results.spread.plot(style='o-',linewidth=2,legend=True)
    ax=results.mavg.plot(style='-',ax=ax,legend=True)
    ax=results.liq_upper.plot(style='--',ax=ax,legend=True)
    ax=results.liq_lower.plot(style='--',ax=ax,legend=True)
    ax=results.sellspread_threshold.plot(style='x-',ax=ax,legend=True)
    ax=results.buyspread_threshold.plot(style='x-',ax=ax,legend=True)
    ax.plot(results.ix[results.liquidated].index, results.spread[results.liquidated],'8', markersize=10, color='r')

    #ax=results.mstd.plot(style='o:',ax=ax)
    plt.grid(True)
    #plt.ylabel('Spread=({}-{})'.format(_stocks[0],_stocks[1]), color="r")
    cursor3 = Cursor(ax4, useblit=True, color='red', linewidth=1)

    ax5 = plt.subplot(224, sharex=ax1)
    #data['spread'].plot(style='o-',marker=r'$\clubsuit$')
    #data['spread'].plot(style='o-')
    plt.setp(ax5.get_xticklabels(), visible=True)
    #results.capital_used.plot(style='o-',linewidth=1,legend=True) #
    results.pnl.plot(style='o-',linewidth=1,legend=True) #
    ax5.plot(results.ix[results.liquidated].index, results.pnl[results.liquidated],'8', markersize=10, color='y')
    #plt.ylabel('cash used', color="b")
    cursor4 = Cursor(ax5, useblit=True, color='red', linewidth=1 )

    #ax5 = plt.subplot(325)
    #results.pnl.plot()
    #plt.ylabel('pnl', color="r")

    #ax5 = plt.subplot(326)
    #results.positions.plot()
    #plt.ylabel('positions', color="r")
    plt.tight_layout()
    plt.subplots_adjust(left=0.1, bottom=0.1, wspace=0.09,hspace=0.09)

    pngfile='{s0}/{s1}-{s2}-{s3}.png'.format(s0=outputdir,s1=_stocks[0],s2=_stocks[1],s3=ymd)
    plt.savefig(pngfile)

    plt.show()



#data = None
#def backtest(_stocks, window=22*12*4, mavg_window_len=22*2, thresholds=(0.4,0.5,0.9,0.8), basic_unit=20, display=False, printout=False):
def backtest(f, tradelog, pda, thresholds=(0.4,0.5,0.9,0.8), basic_unit=20, display=False, printout=False):
    pairtrade = Pairtrade(pda.stocks,thresholds=thresholds,basic_unit = basic_unit,mavg_window_len=pda.mavg_window_len)
    results = pairtrade.run(pda.data)
    if pairtrade.invested:
        pairtrade.liquidate()

    totalreturn=(results.portfolio_value[-1]-pairtrade.capital_base)/pairtrade.capital_base

    accumulativePNL = 0
    totalPNL = 0

    j=0
    for index, row in results.iterrows():
        tnxs= row['transactions']
        tradelog.write("[band][sp={}][{}]({},{},{},{})[pnl={}]\n".format(row['spread'],index.date().strftime('%Y-%m-%d'),row.sellspread_threshold,row.liq_upper,row.liq_lower,row.buyspread_threshold,row.pnl))
        accumulativePNL += row.pnl
        totalPNL += row.pnl
        if tnxs:
            if row['liquidated']:
                if printout:
                    tradelog.write('LIQ:\n') #liquidation
            for txn in tnxs:
                if printout:
                    tradelog.write("[{}][{}][comm:{}]spread:{},amount:{},sid:{},price:{},dt:{},mavg:{},({},{},{},{}),pnl:{},cost:{},mstd:{}\n".\
                    format(j,row.tradedaysincelastliquidation,txn['commission'],row['spread'],txn['amount'],txn['sid'],txn['price'],txn['dt'],\
                           row['mavg'],row.sellspread_threshold,row.liq_upper,row.liq_lower,row.buyspread_threshold,row.pnl,row.cost,row.mstd))
                #print "\tpostion:{}".format(row['positions'])
                j=j+1
            if row['liquidated']:
                if printout:
                    tradelog.write("RealizedPNL:"+str(accumulativePNL) + "\n")
                accumulativePNL = 0
    if display:
        try:
            pairplot(results,pda.data,_stocks,totalreturn,maxdrawdown=min(results.pnl),thresholds=thresholds)
        except:pass

    tradelog.write(("totalPNL:{} total return:{}%, maxdrawdown={}\n".format(totalPNL, totalreturn*100, min(results.pnl))))
    # write to file
    row = results.tail(n=1)

    s="{s1},{s2},{t1},{t2},{t3},{t4},{basic_unit}\n".format(s1=_stocks[0],
                                                            s2=_stocks[1],
                                                            t1=row.sellspread_threshold.values[0],
                                                            t2=row.liq_upper.values[0],
                                                            t3=row.liq_lower.values[0],
                                                            t4=row.buyspread_threshold.values[0],
                                                            basic_unit=basic_unit)
    f.write(s)


    return (totalreturn*100,min(results.pnl))

class pairdata:
    def __init__(self, _stocks, mavg_window_len=22*2):
        end = datetime.today()
        start = end-timedelta(days=22*12*3) #TODO - what if there is no data for one/two stocks 3年的数据比较好
        self.stocks = _stocks
        self.window_len = 0
        self.mavg_window_len = mavg_window_len
        try:
            self.data = load_from_yahoo(stocks=_stocks, indexes={}, start=start, end=end)
            self.data['spread']=self.data[_stocks[0]] - self.data[_stocks[1]]
            self.data = self.data[np.isfinite(self.data['spread'])]
            self.window_len = len(self.data)
            if(self.window_len/self.mavg_window_len<10):
                self.mavg_window_len=self.window_len/10
        except:
            pass

if __name__ == '__main__':
    #ss=[['HMIN','HTHT'],['ATHM','BITA'], ['GSI','SUTR'], ['QUNR','CTRP'], ['ZNH','CEA'], ['SOHU','SINA'], ['SOHU','NTES']]
    ss=[['HMIN','HTHT',50],['ATHM','BITA',50]] #['GSI','SUTR'],
    #ss=[['ZNH','CEA'],['ATHM','BITA']]
    #ss=[['NTES','SOHU']]
    #ss=[['ATHM','BITA']]
    f = open('{}/twintowers-{}.lst'.format(outputdir,ymd),'w')#'a+')
    for _stocks in ss:
        print _stocks
        #pda = pairdata(_stocks, mavg_window_len=22*2)

        basic_unit = _stocks[2]
        _stocks = _stocks[:2]
        tradelog = open('{}/sim-trade-{}-{}-{}.lst'.format(outputdir,_stocks[0],_stocks[1],ymd),'w')
        pda = pairdata(_stocks, mavg_window_len=22*2)
        #backtest(_stocks, window=22*12*3, thresholds=(0.4,0.5,0.9,0.8), display=False,printout=False)
        #window_len=22*12*3
        #mavg_window_len=22*2
        #backtest(_stocks, window=window_len, mavg_window_len=mavg_window_len, thresholds=(0.3,0.3,0.6,0.6), basic_unit=40, display=False, printout=False)

        #mn=pda.data.spread.abs().mean()
        #basic_unit = (seedcap/(10*mn))//10 * 10

        tradelog.write("basic_unit = %d, spread.abs().mean()=%f\n" % (basic_unit,pda.data.spread.abs().mean()))
        r =backtest(f, tradelog, pda, thresholds=(0.3,0.3,1.,1.), basic_unit=basic_unit, display=True, printout=True)
        print r
        continue
        sys.exit(0)

        f2 = open('{}/{}-{}-param.lst'.format(outputdir,_stocks[0],_stocks[1]),'w')
        for i in np.arange(0.1,1,0.2):
            for j in np.arange(0.1,1,0.2):
                thresholds = (i,i,i+j,i+j)
                #print "thresholds", thresholds
                r=backtest(_stocks, window=22*12*3, thresholds=thresholds, basic_unit=20, display=False,printout=False)
                print thresholds,',',r[0],',',r[1]
                line='{},{},{},{},{},{},{}\n'.format(thresholds[0],thresholds[1],thresholds[2],thresholds[3],r[0],r[1],-r[0]/r[1])
                f2.write(line)
        f2.close()
    f.close()

""">>> for i in results.columns:print i
liq_upper
liq_lower
sellspread_threshold
buyspread_threshold
buy
*capital_used
cost
*ending_cash
ending_value
liquidated
mavg
mstd
orders
period_close
period_open
pnl
portfolio_value # mark to market
positions
returns
sell
spread
*starting_cash
starting_value
stock0
tradedaysincelastliquidation
transactions"""